/*

  OpenCL implementation of the FX correlator

 */

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <math.h>
#include <complex.h>
#include <cuComplex.h>
#include <limits.h>
#include<string.h>

#define TILE_WIDTH 16
#define TILE_HEIGHT 16
#define NPOL 2
#define NSTATION 512
#define NTIME 100
#define TOL 1e-3

// All OpenCL headers
#if defined (__APPLE__) || defined(MACOSX)
    #include <OpenCL/opencl.h>
#else
    #include <CL/opencl.h>
#endif 

// name of the source file
const char *simpleSource = "simple_triangle.cl";

// memory pointers on host machine
float complex *array_h;
float complex *product_h;
float complex *compare_h;

// memory pointers on device
cl_mem array_d;
cl_mem product_d;

// OpenCL variables
cl_context cxGPUContext;
cl_command_queue cqCommandQueue;
cl_platform_id cpPlatform;
cl_device_id cdDevice;
cl_program cpProgram;
cl_kernel ckKernel;
size_t szGlobalWorkSize[2]; // total number of threads
size_t szLocalWorkSize[2]; // threads per work group
size_t szParamDataBytes; // byte size of context
size_t szKernelLength; // byte size of kernel
cl_int ciErr1, ciErr2;
char *cPathAndName = NULL;
char *cSourceCL = NULL;

char* oclLoadProgSource(const char* cFilename, const char* cPreamble, size_t* szFinalLength);

void random_complex(float complex* random_num, int length) {
  float a,b;
  //complex<float> random_num[length];
  for(int i=0; i<length; i++){
    a = ((rand()-INT_MAX/2) / (float)(INT_MAX/2)) + 10.0f;
    b = ((rand()-INT_MAX/2) / (float)(INT_MAX/2)) + 10.0f;
    random_num[i] = a + I*b;
  }
}

void outerProduct(float complex *compare_h, float complex *array_h) {
  for(int station1=0; station1<NSTATION; station1++){
    for(int station2=0; station2<=station1; station2++){
      for (int pol1 = 0; pol1<NPOL; pol1++) {
        for (int pol2 = 0; pol2<NPOL; pol2++) {
          complex float sum = 0 + I*0;
          for (int t=0; t<NTIME; t++) {
            sum += array_h[t*NSTATION*NPOL + station1*NPOL + pol1] * ~array_h[t*NSTATION*NPOL + station2*NPOL + pol2];
          }
          compare_h[(( station1*(station1+1)/2 + station2 )*NPOL + pol1 )*NPOL + pol2] = sum;
        }
      }
    }
  }
}

//check that GPU caluculation matches the CPU one
void checkResult(float complex *gpu, float complex *cpu, char *kernel) {

  printf("Checking %s kernel...\n", kernel); fflush(stdout);
  
  for(int i=0; i<NSTATION; i++){
    for (int j=0; j<=i; j++) {
      int k = i*(i+1)/2 + j;
      for (int pol1=0; pol1<NPOL; pol1++) {
	for (int pol2=0; pol2<NPOL; pol2++) {
	  int index = (k*NPOL+pol1)*NPOL+pol2;
	  //printf("%d %d %d %d %d %d\n",i,j,k,pol1,pol2,index);
	  //printf("%f  %f\n", __real__ cpu[index], __imag__ cpu[index]);
	  //printf("%f  %f\n", __real__ gpu[index], __imag__ gpu[index]);
	  assert(cabsf(cpu[index] - gpu[index]) / cabsf(cpu[index]) < TOL);
	}
      }
    }
  }
  //if program makes it this far, then memory was copied successfully
  printf("Outer product operation successful for %s\n\n", kernel);

}

int main(int argc, char **argv) {

  int vecLength = NTIME * NSTATION * NPOL;
  int matLength = (NSTATION+1)*(NSTATION/2)*NPOL*NPOL;
  int numBlocksX = NSTATION/TILE_WIDTH;
  int numBlocksY = NSTATION/TILE_HEIGHT;

  // set global and local dimensions
  szLocalWorkSize[0] = TILE_WIDTH;
  szLocalWorkSize[1] = TILE_WIDTH;
  szGlobalWorkSize[0] = NSTATION;
  szGlobalWorkSize[1] = NSTATION;

  //perform memory allocation
  array_h = (float complex *) malloc(vecLength*sizeof(float complex));
  product_h = (float complex *) malloc(matLength*sizeof(float complex));
  compare_h = (float complex *) malloc(matLength*sizeof(float complex));

  // Get the platforms
  ciErr1 = clGetPlatformIDs(1, &cpPlatform, NULL);
  if (ciErr1 != CL_SUCCESS) {
    printf("PlatformID failed\n");
    exit(0);
  }

  // Get the devices
  ciErr1 = clGetDeviceIDs(cpPlatform, CL_DEVICE_TYPE_GPU, 1, &cdDevice, NULL);
  if (ciErr1 != CL_SUCCESS) {
    printf("DeviceID failed\n");
    exit(0);
  }

  //Create the context
  cxGPUContext = clCreateContext(0, 1, &cdDevice, NULL, NULL, &ciErr1);
  if (ciErr1 != CL_SUCCESS) {
    printf("Context failed\n");
    exit(0);
  }
  
  // Create a command-queue
  cqCommandQueue = clCreateCommandQueue(cxGPUContext, cdDevice, 0, &ciErr1);
  if (ciErr1 != CL_SUCCESS) {
    printf("CommandQueue failed\n");
    exit(0);
  }

  // Allocate the OpenCL buffer memory objects for source and result on the device GMEM
  array_d = clCreateBuffer(cxGPUContext, CL_MEM_READ_ONLY, sizeof(float complex) * vecLength, NULL, &ciErr1);
  ciErr1 |= ciErr2;
  product_d = clCreateBuffer(cxGPUContext, CL_MEM_WRITE_ONLY, sizeof(float complex) * matLength, NULL, &ciErr2);
  ciErr1 |= ciErr2;
  if (ciErr1 != CL_SUCCESS) {
    printf("Error in clCreateBuffer, Line %u in file %s !!!\n\n", __LINE__, __FILE__);
    exit(0);
  }

  // Fill input array with random numbers
  random_complex(array_h, vecLength);

  // Do host computation
  outerProduct(compare_h, array_h);

  // Read the OpenCL kernel in from source file
  //cPathAndName = shrFindFilePath(simpleSource, argv[0]);
  cSourceCL = oclLoadProgSource(simpleSource, "", &szKernelLength);  
  
  // Create the program
  cpProgram = clCreateProgramWithSource(cxGPUContext, 1, (const char **)&cSourceCL, &szKernelLength, &ciErr1);
  printf("%s\n", cSourceCL);
  if (ciErr1 != CL_SUCCESS) {
    printf("Error in clCreateProgramWithSource, Line %u in file %s !!!\n\n", __LINE__, __FILE__);
    exit(0);
  }

  // compile flags
  char* flags = "-cl-fast-relaxed-math";

  // compile the kernel
  ciErr1 = clBuildProgram(cpProgram, 0, NULL, NULL, NULL, NULL);
  if (ciErr1 != CL_SUCCESS) {
    // write out the build log, then exit
    char cBuildLog[10240];
    clGetProgramBuildInfo(cpProgram, cdDevice, CL_PROGRAM_BUILD_LOG, 
                          sizeof(cBuildLog), cBuildLog, NULL );
    printf("\n\nBuild Log:\n%s\n\n", cBuildLog);
    printf("Error in clBuildProgram, Line %u in file %s !!!\n\n", __LINE__, __FILE__);
    exit(0);
  }

  // Create the kernel
  ckKernel = clCreateKernel(cpProgram, "SimpleTriangle", &ciErr1);
  printf("clCreateKernel (SimpleTriangle)...\n"); 
  if (ciErr1 != CL_SUCCESS) {
    printf("Error in clCreateKernel, Line %u in file %s !!!\n\n", __LINE__, __FILE__);
    exit(0);
  }

  // Set the Argument values
  ciErr1 = clSetKernelArg(ckKernel, 0, sizeof(cl_mem), (void*)&array_d);
  if (ciErr1 != CL_SUCCESS){
    printf("Error %d in clSetKernelArg, Line %u in file %s !!! %d %d %d %d %d %d\n\n", ciErr1, __LINE__, __FILE__,
	   CL_INVALID_KERNEL, CL_INVALID_ARG_INDEX, CL_INVALID_ARG_VALUE, 
	   CL_INVALID_MEM_OBJECT, CL_INVALID_SAMPLER, CL_INVALID_ARG_SIZE);
    exit(0);
  }
  ciErr1 |= clSetKernelArg(ckKernel, 1, sizeof(cl_mem), (void*)&product_d);
  if (ciErr1 != CL_SUCCESS){
    printf("Error %d in clSetKernelArg, Line %u in file %s !!!\n\n", ciErr1, __LINE__, __FILE__);
    exit(0);
  }
  
  
  // Asynchronous write of data to GPU device
  ciErr1 = clEnqueueWriteBuffer(cqCommandQueue, array_d, CL_FALSE, 0, sizeof(float complex) * vecLength, 
				array_h, 0, NULL, NULL);
  if (ciErr1 != CL_SUCCESS) {
    printf("Error in clEnqueueWriteBuffer, Line %u in file %s !!!\n\n", __LINE__, __FILE__);
    exit(0);
  }

  // Launch kernel
  ciErr1 = clEnqueueNDRangeKernel(cqCommandQueue, ckKernel, 2, NULL, szGlobalWorkSize, szLocalWorkSize, 0, NULL, NULL);
  if (ciErr1 != CL_SUCCESS) {
    printf("Error in clEnqueueNDRangeKernel, Line %u in file %s !!!\n\n", __LINE__, __FILE__);
    exit(0);
  }

  // Synchronous/blocking read of results, and check accumulated errors
  ciErr1 = clEnqueueReadBuffer(cqCommandQueue, product_d, CL_TRUE, 0, sizeof(float complex) * matLength, 
			       product_h, 0, NULL, NULL);
  if (ciErr1 != CL_SUCCESS) {
    printf("Error in clEnqueueReadBuffer, Line %u in file %s !!!\n\n", __LINE__, __FILE__);
    exit(0);
  }

  // Compute and compare results for golden-host and report errors and pass/fail
  printf("Comparing against Host/C++ computation...\n\n"); 
  checkResult(product_h, compare_h, "SimpleTriangle");

  clReleaseCommandQueue(cqCommandQueue);
  clReleaseContext(cxGPUContext);

  free(array_h);
  free(product_h);
  free(compare_h);

}


char* oclLoadProgSource(const char* cFilename, const char* cPreamble, size_t* szFinalLength)
{
    // locals 
    FILE* pFileStream = NULL;
    size_t szSourceLength;

    // open the OpenCL source code file
    #ifdef _WIN32   // Windows version
        if(fopen_s(&pFileStream, cFilename, "rb") != 0) 
        {       
            return NULL;
        }
    #else           // Linux version
        pFileStream = fopen(cFilename, "rb");
        if(pFileStream == 0) 
        {       
            return NULL;
        }
    #endif

    size_t szPreambleLength = strlen(cPreamble);

    // get the length of the source code
    fseek(pFileStream, 0, SEEK_END); 
    szSourceLength = ftell(pFileStream);
    fseek(pFileStream, 0, SEEK_SET); 

    // allocate a buffer for the source code string and read it in
    char* cSourceString = (char *)malloc(szSourceLength + szPreambleLength + 1); 
    memcpy(cSourceString, cPreamble, szPreambleLength);
    if (fread((cSourceString) + szPreambleLength, szSourceLength, 1, pFileStream) != 1)
    {
        fclose(pFileStream);
        free(cSourceString);
        return 0;
    }

    // close the file and return the total length of the combined (preamble + source) string
    fclose(pFileStream);
    if(szFinalLength != 0)
    {
        *szFinalLength = szSourceLength + szPreambleLength;
    }
    cSourceString[szSourceLength + szPreambleLength] = '\0';

    return cSourceString;
}
